#include "population.hh"
#include "constants.hh"
#include <algorithm>
#include <map>
#include <iostream>
using namespace std;
using namespace evolution;

void
Population::add_individual(individual_t individual) {
    current_generation.push_back(individual);
    place_in_species(individual, this->species);
}

void
Population::calculate_fitness() {
    for (auto individual: current_generation) individual->calculate_fitness();
}


void
Population::next_generation() {
    std::map<individual_t, float> species_fitnesses;

    float total_fitness = 0;
    float tmp = 0;
    for (auto& species: this->species) {        
        species.second.sort (
            [](individual_t const& a, individual_t const& b) {
                return a->get_fitness() > b->get_fitness();
            });

        float species_fitness = 0;
        //cout << "There are " << species.second.size() << " individuals in a species!\n";
        for (individual_t individual: species.second) {
            float fitness = individual->get_fitness();
            //cout << "Individual fitness = " << fitness << endl;
            if (fitness > tmp) tmp = fitness;
            species_fitness += fitness;
        }
        total_fitness += species_fitness / species.second.size();
        
        species_fitnesses[species.first] = species_fitness / species.second.size();
        //cout << "adjusted fitness = " << species_fitnesses[species.first] << endl;
    }

    std::vector<individual_t> next_generation;
    int i = 0;
    for (auto& species: this->species) {
        const float fitness = species_fitnesses[species.first];
        const int old_num = species.second.size();

        if (fitness < 0.9f*total_fitness/species_count()) continue;
        int new_num = (current_generation.size()) * (fitness / total_fitness) - 1;


        individual_t best = get_best_individual_in(species.second);
        cout << "species " << i++ << " (members: " << species.second.size()
             << ") gets: " << new_num << " offspring!\n";
        cout << "\tfitness: " << fitness << " (total fitness: " << total_fitness << ")\n";
        if (new_num > 0) next_generation.push_back(best->twin());      
        
        for (int i = 1; i < new_num; ++i) {      
            next_generation.push_back(best->clone());
            if (species.second.size() > old_num*0.8f && species.second.size() > 1) {
                species.second.remove(best);                
                best = get_best_individual_in(species.second);
            }
        }
    }

    individual_t best = get_best_individual();
    for (unsigned i = next_generation.size(); i < current_generation.size(); ++i) {
        next_generation.push_back(best->clone());
    }

    cout << "total individuals in new generation: " << next_generation.size() << endl;
    cout << "best individual in current generation: " << tmp << endl;

    species_t next;
    for (auto const& individual: next_generation) {
        place_in_species(individual, next);
    }

    species = next;
    current_generation = next_generation;
}

void
Population::place_in_species (individual_t individual, species_t& species_) {
    for (auto& species: species_) {
        Individual* prototype = species.first.get();
        std::list<individual_t>& list = species.second;

        if (individual->genetic_distance(prototype)
            < SPECIATION_THRESHOLD) {
            
            list.push_back(individual);
            return;
        }
    }

    species_.push_back(make_pair(individual, std::list<individual_t>{individual}));
}









    
